{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6D0GXH9w4td"
      },
      "source": [
        "# Train GRBMs on GMM, MNIST, CelebA, and CIFAR10 \n",
        "\n",
        "Copyright 2022 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "\n",
        "Required packages:\n",
        "\n",
        "*   pytorch_fid\n",
        "*   seaborn\n",
        "*   pandas\n",
        "*   matplotlib\n",
        "\n",
        "To change the hyperparameters, please go to the config function.\n",
        "\n",
        "Important hyperparameters include: \n",
        "\n",
        "*   CD_step\n",
        "*   hidden_size\n",
        "*   inference_method: ['Gibbs', 'Langevin', 'Langevin-Gibbs'] "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iTdKNOxXccgO"
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function\n",
        "import logging\n",
        "from collections import defaultdict\n",
        "import seaborn as sns\n",
        "import pandas as pd\n",
        "import os\n",
        "import pdb\n",
        "import PIL\n",
        "import imageio\n",
        "import math\n",
        "import json\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms, utils\n",
        "from torchvision.datasets.utils import download_file_from_google_drive, check_integrity, verify_str_arg\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "from PIL import Image  # NOQA\n",
        "import matplotlib\n",
        "from typing import Any, List, Union, Optional, Tuple, Callable\n",
        "from functools import partial\n",
        "from threading import Thread\n",
        "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "\n",
        "matplotlib.use('Agg')\n",
        "import matplotlib.pyplot as plt  # NOQA\n",
        "\n",
        "# Commented out IPython magic to ensure Python compatibility.\n",
        "# %matplotlib inline\n",
        "EPS = 1e-7\n",
        "torch.set_default_dtype(torch.float64)\n",
        "sns.set_theme(style=\"darkgrid\")\n",
        "\"\"\"Here are some helper functions for training and visualization.\"\"\"\n",
        "\n",
        "\n",
        "class TensorEncoder(json.JSONEncoder):\n",
        "\n",
        "  def default(self, obj):\n",
        "    if isinstance(obj, torch.Tensor):\n",
        "      return obj.cpu().numpy().tolist()\n",
        "    return json.JSONEncoder.default(self, obj)\n",
        "\n",
        "\n",
        "def setup_logging(log_level, log_file, logger_name=\"exp_logger\"):\n",
        "  \"\"\" Setup logging \"\"\"\n",
        "  numeric_level = getattr(logging, log_level.upper(), None)\n",
        "  if not isinstance(numeric_level, int):\n",
        "    raise ValueError(\"Invalid log level: %s\" % log_level)\n",
        "\n",
        "  logging.basicConfig(\n",
        "      filename=log_file,\n",
        "      filemode=\"w\",\n",
        "      format=\n",
        "      \"%(levelname)-5s | %(asctime)s | File %(filename)-20s | Line %(lineno)-5d | %(message)s\",\n",
        "      datefmt=\"%m/%d/%Y %I:%M:%S %p\",\n",
        "      level=numeric_level)\n",
        "\n",
        "  # define a Handler which writes messages to the sys.stderr\n",
        "  console = logging.StreamHandler()\n",
        "  console.setLevel(numeric_level)\n",
        "  # set a format which is simpler for console use\n",
        "  formatter = logging.Formatter(\n",
        "      \"%(levelname)-5s | %(asctime)s | %(filename)-25s | line %(lineno)-5d: %(message)s\"\n",
        "  )\n",
        "  # tell the handler to use this format\n",
        "  console.setFormatter(formatter)\n",
        "  # add the handler to the root logger\n",
        "  logging.getLogger(logger_name).addHandler(console)\n",
        "\n",
        "  return get_logger(logger_name)\n",
        "\n",
        "\n",
        "def get_logger(logger_name=\"exp_logger\"):\n",
        "  return logging.getLogger(logger_name)\n",
        "\n",
        "\n",
        "def cosine_schedule(eta_min=0, eta_max=1, T=10):\n",
        "  return [\n",
        "      eta_min + (eta_max - eta_min) * (1 + math.cos(tt * math.pi / T)) / 2\n",
        "      for tt in range(T)\n",
        "  ]\n",
        "\n",
        "\n",
        "def fig2img(fig):\n",
        "  import io\n",
        "  buf = io.BytesIO()\n",
        "  fig.savefig(buf, bbox_inches='tight')\n",
        "  buf.seek(0)\n",
        "  img = Image.open(buf)\n",
        "  return img\n",
        "\n",
        "\n",
        "def show_img(matrix, title):\n",
        "  plt.figure()\n",
        "  plt.axis('off')\n",
        "  plt.gray()\n",
        "  img = np.array(matrix, np.float64)\n",
        "  # plt.imshow(img, interpolation='bilinear')\n",
        "  plt.imshow(img)\n",
        "  plt.title(title)\n",
        "\n",
        "  fig = plt.gcf()\n",
        "  img_out = fig2img(fig)\n",
        "  plt.close()\n",
        "\n",
        "  return img_out\n",
        "\n",
        "\n",
        "def save_gif_fancy(imgs, nrow, save_name):\n",
        "  imgs = (show_img(\n",
        "      utils.make_grid(xx[1],\n",
        "                      nrow=nrow,\n",
        "                      normalize=False,\n",
        "                      padding=1,\n",
        "                      pad_value=1.0).permute(1, 2, 0).cpu().numpy(),\n",
        "      f'sample at {xx[0]:03d} step') for xx in imgs)\n",
        "  img = next(imgs)\n",
        "  img.save(fp=save_name,\n",
        "           format='GIF',\n",
        "           append_images=imgs,\n",
        "           save_all=True,\n",
        "           duration=400,\n",
        "           loop=0)\n",
        "\n",
        "\n",
        "def save(model, results_folder, epoch):\n",
        "  data = {'epoch': epoch, 'model': model.state_dict()}\n",
        "  torch.save(data, f'{results_folder}/model-{epoch}.pt')\n",
        "\n",
        "\n",
        "def load(model, results_folder, epoch):\n",
        "  data = torch.load(f'{results_folder}/model-{epoch}.pt')\n",
        "  model.load_state_dict(data['model'])\n",
        "\n",
        "\n",
        "def train(model,\n",
        "          train_loader,\n",
        "          optimizer,\n",
        "          config,\n",
        "          is_anneal_data_noise=False,\n",
        "          std=0.0):\n",
        "  model.train()\n",
        "  for ii, (data, _) in enumerate(tqdm(train_loader)):\n",
        "    if config['cuda']:\n",
        "      data = data.cuda()\n",
        "\n",
        "    if is_anneal_data_noise:\n",
        "      data += torch.randn_like(data) * std\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "    model.CD_grad(data)\n",
        "    if config['clip_norm'] \u003e 0:\n",
        "      nn.utils.clip_grad_norm_(model.parameters(), config['clip_norm'])\n",
        "    optimizer.step()\n",
        "\n",
        "    if ii == len(train_loader) - 1:\n",
        "      recon_loss = model.reconstruction(data).item()\n",
        "\n",
        "  return recon_loss\n",
        "\n",
        "\n",
        "def unnormalize_img_tuple(img_tuple, mean, std):\n",
        "  if isinstance(std, torch.Tensor):\n",
        "    mean = mean.view(1, -1, 1, 1).to(img_tuple[0][1].device)\n",
        "    std = std.view(1, -1, 1, 1).to(img_tuple[0][1].device)\n",
        "\n",
        "  return [(xx[0], (xx[1] * std + mean).clamp(min=0, max=1)) for xx in img_tuple]\n",
        "\n",
        "\n",
        "def visualize_sampling(model, epoch, config, tag=None, is_show_gif=True):\n",
        "  tag = '' if tag is None else tag\n",
        "  B, C, H, W = config['sampling_batch_size'], config['channel'], config[\n",
        "      'height'], config['width']\n",
        "  v_init = torch.randn(B, C, H, W).cuda() * config['neg_MC_init_var']\n",
        "  v_list = model.sampling(v_init,\n",
        "                          num_steps=config['sampling_steps'],\n",
        "                          save_gap=config['sampling_gap'])\n",
        "\n",
        "  if config['dataset'] == 'GMM':\n",
        "    samples = v_list[-1][1].view(B, -1).cpu().numpy()\n",
        "    vis_2D_samples(samples, config, tags=f'{epoch:05d}')\n",
        "    vis_density_GRBM(model, config, epoch=epoch)\n",
        "  else:\n",
        "    if is_show_gif:\n",
        "      v_list = unnormalize_img_tuple(v_list, config['img_mean'],\n",
        "                                     config['img_std'])\n",
        "      save_gif_fancy(\n",
        "          v_list, config['sampling_nrow'],\n",
        "          f\"{config['exp_folder']}/sample_imgs_epoch_{epoch:05d}{tag}.gif\")\n",
        "      img_vis = v_list[-1][1]\n",
        "    else:\n",
        "      if isinstance(config['img_std'], torch.Tensor):\n",
        "        mean = config['img_mean'].view(1, -1, 1, 1).cuda()\n",
        "        std = config['img_std'].view(1, -1, 1, 1).cuda()\n",
        "      else:\n",
        "        mean = config['img_mean']\n",
        "        std = config['img_std']\n",
        "\n",
        "      img_vis = (v_list[-1][1] * std + mean).clamp(min=0, max=1)\n",
        "\n",
        "    utils.save_image(\n",
        "        utils.make_grid(img_vis,\n",
        "                        nrow=config['sampling_nrow'],\n",
        "                        normalize=False,\n",
        "                        padding=1,\n",
        "                        pad_value=1.0).cpu(),\n",
        "        f\"{config['exp_folder']}/sample_imgs_epoch_{epoch:05d}{tag}.png\")\n",
        "\n",
        "\n",
        "def vis_2D_samples(samples, config, tags=None):\n",
        "  f, ax = plt.subplots(figsize=(6, 6))\n",
        "  # sns.scatterplot(x=samples[:, 0], y=samples[:, 1], color=\"#4CB391\", s=5)\n",
        "  sns.scatterplot(x=samples[:, 0], y=samples[:, 1], color=\"#4CB391\")\n",
        "  # sns.histplot(x=samples[:, 0], y=samples[:, 1],\n",
        "  #              bins=50, pthresh=.1, cmap=\"mako\")\n",
        "  # sns.kdeplot(x=samples[:, 0], y=samples[:, 1],\n",
        "  #             levels=5, cmap=\"crest\",\n",
        "  #             fill=True,\n",
        "  #             alpha=0.5,\n",
        "  #             cut=2)\n",
        "  # sns.kdeplot(x=samples[:, 0],\n",
        "  #             y=samples[:, 1],\n",
        "  #             levels=5,\n",
        "  #             color=\"b\",\n",
        "  #             linewidths=1)\n",
        "  ax.set(xlim=(-10, 10))\n",
        "  ax.set(ylim=(-10, 10))\n",
        "  plt.show()\n",
        "  plt.savefig(f\"{config['exp_folder']}/samples_{tags}.png\", bbox_inches='tight')\n",
        "  plt.close()\n",
        "\n",
        "\n",
        "def vis_density_GMM(model, config):\n",
        "  fig, ax = plt.subplots()\n",
        "  x_density, y_density = 500, 500\n",
        "  xses = np.linspace(-10, 10, x_density)\n",
        "  yses = np.linspace(-10, 10, y_density)\n",
        "  xy = torch.tensor([[[x, y] for x in xses] for y in yses]).view(-1, 2).cuda()\n",
        "  log_density_values = model.log_prob(xy)\n",
        "  log_density_values = log_density_values.detach().view(\n",
        "      x_density, y_density).cpu().numpy()\n",
        "  dx = (xses[1] - xses[0]) / 2\n",
        "  dy = (yses[1] - yses[0]) / 2\n",
        "  extent = [xses[0] - dx, xses[-1] + dx, yses[0] - dy, yses[-1] + dy]\n",
        "  im = ax.imshow(np.exp(log_density_values),\n",
        "                 extent=extent,\n",
        "                 origin='lower',\n",
        "                 cmap='viridis')\n",
        "  ax.set_xlabel('x')\n",
        "  ax.set_ylabel('y')\n",
        "  divider = make_axes_locatable(ax)\n",
        "  cax = divider.append_axes('right', size='5%', pad=0.05)\n",
        "  cb = fig.colorbar(im, cax=cax)\n",
        "  cb.set_label('probability density')\n",
        "  plt.show()\n",
        "  plt.savefig(f\"{config['exp_folder']}/GMM_density.png\", bbox_inches='tight')\n",
        "  plt.close()\n",
        "\n",
        "\n",
        "def vis_density_GRBM(model, config, epoch=None):\n",
        "  fig, ax = plt.subplots()\n",
        "  x_density, y_density = 500, 500\n",
        "  xses = np.linspace(-10, 10, x_density)\n",
        "  yses = np.linspace(-10, 10, y_density)\n",
        "  xy = torch.tensor([[[x, y] for x in xses] for y in yses]).view(-1, 2).cuda()\n",
        "  eng_val = -model.marginal_energy(xy)\n",
        "  eng_val = eng_val.detach().view(x_density, y_density).cpu().numpy()\n",
        "  dx = (xses[1] - xses[0]) / 2\n",
        "  dy = (yses[1] - yses[0]) / 2\n",
        "  extent = [xses[0] - dx, xses[-1] + dx, yses[0] - dy, yses[-1] + dy]\n",
        "  im = ax.imshow(eng_val, extent=extent, origin='lower', cmap='viridis')\n",
        "  ax.set_xlabel('x')\n",
        "  ax.set_ylabel('y')\n",
        "  divider = make_axes_locatable(ax)\n",
        "  cax = divider.append_axes('right', size='5%', pad=0.05)\n",
        "  cb = fig.colorbar(im, cax=cax)\n",
        "  cb.set_label('negative energy')\n",
        "  plt.show()\n",
        "  plt.savefig(f\"{config['exp_folder']}/GRBM_density_{epoch:05d}.png\",\n",
        "              bbox_inches='tight')\n",
        "  plt.close()\n",
        "\n",
        "\n",
        "def plot_curves(stats_dict, config):\n",
        "  df = pd.DataFrame(data=stats_dict)\n",
        "\n",
        "  # visualize variance\n",
        "  sns_plot = sns.lineplot(data=df, x=f'Epoch', y=f'Variance')\n",
        "  fig = sns_plot.get_figure()\n",
        "  fig.savefig(\n",
        "      f\"{config['exp_folder']}/Var_curve_{config['resume']:05d}_{config['epochs']:05d}.png\",\n",
        "      bbox_inches='tight')\n",
        "  plt.clf()\n",
        "\n",
        "  sns_plot = sns.lineplot(data=df, x=f'Epoch', y=f'Log Variance')\n",
        "  fig = sns_plot.get_figure()\n",
        "  fig.savefig(\n",
        "      f\"{config['exp_folder']}/Log_var_curve_{config['resume']:05d}_{config['epochs']:05d}.png\",\n",
        "      bbox_inches='tight')\n",
        "  plt.clf()\n",
        "\n",
        "  # visualize reconstruction loss\n",
        "  sns_plot = sns.lineplot(data=df, x=f'Epoch', y=f'Reconstruction MSE Loss')\n",
        "  fig = sns_plot.get_figure()\n",
        "  fig.savefig(\n",
        "      f\"{config['exp_folder']}/Recon_curve_{config['resume']:05d}_{config['epochs']:05d}.png\",\n",
        "      bbox_inches='tight')\n",
        "  plt.clf()\n",
        "\n",
        "\n",
        "class ReplayBuffer(object):\n",
        "\n",
        "  def __init__(self, buffer_size=5) -\u003e None:\n",
        "    self.buffer_size = buffer_size\n",
        "    self._queue = []\n",
        "\n",
        "  def enqueue(self, x):\n",
        "    if len(self._queue) \u003c self.buffer_size:\n",
        "      self._queue += [x]\n",
        "    else:\n",
        "      self._queue.pop(0)\n",
        "      self._queue += [x]\n",
        "\n",
        "  def empty(self):\n",
        "    return len(self._queue) == 0\n",
        "\n",
        "  def sample(self, num_sample=128):\n",
        "    assert len(self._queue) \u003e 0\n",
        "    num_sample_total = sum([qq[0].shape[0] for qq in self._queue])\n",
        "    rand_idx = torch.randperm(num_sample_total)\n",
        "    idx = rand_idx[:num_sample]\n",
        "\n",
        "    h_out = []\n",
        "    for ii in range(len(self._queue[0])):\n",
        "      h_out += [torch.concat([qq[ii] for qq in self._queue], dim=0)[idx, :]]\n",
        "\n",
        "    return h_out\n",
        "\n",
        "\n",
        "\"\"\"Here are some datasets.\"\"\"\n",
        "\n",
        "\n",
        "class GMMDataset(torch.utils.data.Dataset):\n",
        "\n",
        "  def __init__(self, samples):\n",
        "    self.samples = samples\n",
        "\n",
        "  def __len__(self):\n",
        "    return self.samples.shape[0]\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    return self.samples[idx, :], torch.ones(1).to(self.samples.device)\n",
        "\n",
        "\n",
        "class GMM(nn.Module):\n",
        "  \"\"\" Gaussian Mixture Models \n",
        "      N.B.: covariance is assumed to be diagonal\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, w, mu, sigma):\n",
        "    \"\"\"\n",
        "      p(x) = sum_i w[i] N(mu[i], sigma[i]^2 * I)\n",
        "\n",
        "      config:\n",
        "        w: shape K X 1, mixture coefficients, must sum to 1\n",
        "        mu: shape K X D, mean\n",
        "        sigma: shape K X D, (diagonal) variance \n",
        "    \"\"\"\n",
        "    super().__init__()\n",
        "    self.register_buffer('w', w)\n",
        "    self.register_buffer('mu', mu)\n",
        "    self.register_buffer('sigma', sigma)\n",
        "    self.K = w.shape[0]\n",
        "    self.D = mu.shape[1]\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def log_gaussian(self, x, mu, sigma):\n",
        "    \"\"\" log density of single (diagonal-covariance) multivariate Gaussian\"\"\"\n",
        "    return -0.5 * ((x - mu)**2 / sigma**2).sum(dim=1) - 0.5 * (\n",
        "        self.D * np.log(2 * np.pi) + torch.log(torch.prod(sigma**2)))\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def log_prob(self, x):\n",
        "    return torch.logsumexp(\n",
        "        torch.stack([\n",
        "            torch.log(self.w[kk]) +\n",
        "            self.log_gaussian(x, self.mu[kk], self.sigma[kk])\n",
        "            for kk in range(self.K)\n",
        "        ]), 0)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def sampling(self, num_samples):\n",
        "    m = torch.distributions.Categorical(self.w)\n",
        "    idx = m.sample((num_samples,))\n",
        "    return self.mu[idx, :] + torch.randn(num_samples, self.D).to(\n",
        "        self.w.device) * self.sigma[idx, :]\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def langevin_sampling(self, x, num_steps=10, eta=1.0e+0, is_anneal=False):\n",
        "    eta_list = cosine_schedule(eta_max=eta, T=num_steps)\n",
        "    for ii in range(num_steps):\n",
        "      eta_ii = eta_list[ii] if is_anneal else eta\n",
        "      x = x.detach()\n",
        "      x.requires_grad = True\n",
        "      eng = -self.log_prob(x).sum()\n",
        "      grad = torch.autograd.grad(eng, x)[0]\n",
        "      x = x - eta_ii * grad + torch.randn_like(x) * np.sqrt(eta_ii * 2)\n",
        "\n",
        "    return x.detach()\n",
        "\n",
        "\n",
        "class CelebA(datasets.VisionDataset):\n",
        "  \"\"\"`Large-scale CelebFaces Attributes (CelebA) Dataset \u003chttp://mmlab.ie.cuhk.edu.hk/projects/CelebA.html\u003e`_ Dataset.\n",
        "\n",
        "    config:\n",
        "        root (string): Root directory where images are downloaded to.\n",
        "        split (string): One of {'train', 'valid', 'test', 'all'}.\n",
        "            Accordingly dataset is selected.\n",
        "        target_type (string or list, optional): Type of target to use, ``attr``, ``identity``, ``bbox``,\n",
        "            or ``landmarks``. Can also be a list to output a tuple with all specified target types.\n",
        "            The targets represent:\n",
        "                ``attr`` (np.array shape=(40,) dtype=int): binary (0, 1) labels for attributes\n",
        "                ``identity`` (int): label for each person (data points with the same identity are the same person)\n",
        "                ``bbox`` (np.array shape=(4,) dtype=int): bounding box (x, y, width, height)\n",
        "                ``landmarks`` (np.array shape=(10,) dtype=int): landmark points (lefteye_x, lefteye_y, righteye_x,\n",
        "                    righteye_y, nose_x, nose_y, leftmouth_x, leftmouth_y, rightmouth_x, rightmouth_y)\n",
        "            Defaults to ``attr``. If empty, ``None`` will be returned as target.\n",
        "        transform (callable, optional): A function/transform that  takes in an PIL image\n",
        "            and returns a transformed version. E.g, ``transforms.ToTensor``\n",
        "        target_transform (callable, optional): A function/transform that takes in the\n",
        "            target and transforms it.\n",
        "        download (bool, optional): If true, downloads the dataset from the internet and\n",
        "            puts it in root directory. If dataset is already downloaded, it is not\n",
        "            downloaded again.\n",
        "    \"\"\"\n",
        "\n",
        "  base_folder = \"celeba\"\n",
        "  # There currently does not appear to be a easy way to extract 7z in python (without introducing additional\n",
        "  # dependencies). The \"in-the-wild\" (not aligned+cropped) images are only in 7z, so they are not available\n",
        "  # right now.\n",
        "  file_list = [\n",
        "      # File ID                         MD5 Hash                            Filename\n",
        "      (\"0B7EVK8r0v71pZjFTYXZWM3FlRnM\", \"00d2c5bc6d35e252742224ab0c1e8fcb\",\n",
        "       \"img_align_celeba.zip\"),\n",
        "      # (\"0B7EVK8r0v71pbWNEUjJKdDQ3dGc\", \"b6cd7e93bc7a96c2dc33f819aa3ac651\", \"img_align_celeba_png.7z\"),\n",
        "      # (\"0B7EVK8r0v71peklHb0pGdDl6R28\", \"b6cd7e93bc7a96c2dc33f819aa3ac651\", \"img_celeba.7z\"),\n",
        "      (\"0B7EVK8r0v71pblRyaVFSWGxPY0U\", \"75e246fa4810816ffd6ee81facbd244c\",\n",
        "       \"list_attr_celeba.txt\"),\n",
        "      (\"1_ee_0u7vcNLOfNLegJRHmolfH5ICW-XS\", \"32bd1bd63d3c78cd57e08160ec5ed1e2\",\n",
        "       \"identity_CelebA.txt\"),\n",
        "      (\"0B7EVK8r0v71pbThiMVRxWXZ4dU0\", \"00566efa6fedff7a56946cd1c10f1c16\",\n",
        "       \"list_bbox_celeba.txt\"),\n",
        "      (\"0B7EVK8r0v71pd0FJY3Blby1HUTQ\", \"cc24ecafdb5b50baae59b03474781f8c\",\n",
        "       \"list_landmarks_align_celeba.txt\"),\n",
        "      # (\"0B7EVK8r0v71pTzJIdlJWdHczRlU\", \"063ee6ddb681f96bc9ca28c6febb9d1a\", \"list_landmarks_celeba.txt\"),\n",
        "      (\"0B7EVK8r0v71pY0NSMzRuSXJEVkk\", \"d32c9cbf5e040fd4025c592c306e6668\",\n",
        "       \"list_eval_partition.txt\"),\n",
        "  ]\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      root: str,\n",
        "      split: str = \"train\",\n",
        "      target_type: Union[List[str], str] = \"attr\",\n",
        "      transform: Optional[Callable] = None,\n",
        "      target_transform: Optional[Callable] = None,\n",
        "      download: bool = False,\n",
        "  ) -\u003e None:\n",
        "    self.return_imgid = 0\n",
        "    import pandas\n",
        "    super(CelebA, self).__init__(root,\n",
        "                                 transform=transform,\n",
        "                                 target_transform=target_transform)\n",
        "    self.split = split\n",
        "    if isinstance(target_type, list):\n",
        "      self.target_type = target_type\n",
        "    else:\n",
        "      self.target_type = [target_type]\n",
        "\n",
        "    if not self.target_type and self.target_transform is not None:\n",
        "      raise RuntimeError(\n",
        "          'target_transform is specified but target_type is empty')\n",
        "\n",
        "    if download:\n",
        "      self.download()\n",
        "\n",
        "    split_map = {\n",
        "        \"train\": 0,\n",
        "        \"valid\": 1,\n",
        "        \"test\": 2,\n",
        "        \"all\": None,\n",
        "    }\n",
        "    split_ = split_map[verify_str_arg(split.lower(), \"split\",\n",
        "                                      (\"train\", \"valid\", \"test\", \"all\"))]\n",
        "\n",
        "    fn = partial(os.path.join, self.root, self.base_folder)\n",
        "    splits = pandas.read_csv(fn(\"list_eval_partition.txt\"),\n",
        "                             delim_whitespace=True,\n",
        "                             header=None,\n",
        "                             index_col=0)\n",
        "\n",
        "    mask = slice(None) if split_ is None else (splits[1] == split_)\n",
        "\n",
        "    self.filename = splits[mask].index.values\n",
        "\n",
        "  def _check_integrity(self) -\u003e bool:\n",
        "    for (_, md5, filename) in self.file_list:\n",
        "      fpath = os.path.join(self.root, self.base_folder, filename)\n",
        "      _, ext = os.path.splitext(filename)\n",
        "      # Allow original archive to be deleted (zip and 7z)\n",
        "      # Only need the extracted images\n",
        "      if ext not in [\".zip\", \".7z\"] and not check_integrity(fpath, md5):\n",
        "        return False\n",
        "\n",
        "    # Should check a hash of the images\n",
        "    return os.path.isdir(\n",
        "        os.path.join(self.root, self.base_folder, \"img_align_celeba\"))\n",
        "\n",
        "  def download(self) -\u003e None:\n",
        "    import zipfile\n",
        "\n",
        "    if self._check_integrity():\n",
        "      print('Files already downloaded and verified')\n",
        "      return\n",
        "\n",
        "    for (file_id, md5, filename) in self.file_list:\n",
        "      download_file_from_google_drive(file_id,\n",
        "                                      os.path.join(self.root, self.base_folder),\n",
        "                                      filename, md5)\n",
        "\n",
        "    with zipfile.ZipFile(\n",
        "        os.path.join(self.root, self.base_folder, \"img_align_celeba.zip\"),\n",
        "        \"r\") as f:\n",
        "      f.extractall(os.path.join(self.root, self.base_folder))\n",
        "\n",
        "  def __getitem__(self, index: int) -\u003e Tuple[Any, Any]:\n",
        "    X = PIL.Image.open(\n",
        "        os.path.join(self.root, self.base_folder, \"img_align_celeba\",\n",
        "                     self.filename[index]))\n",
        "\n",
        "    if self.transform is not None:\n",
        "      X = self.transform(X)\n",
        "\n",
        "    target = torch.zeros(X.shape[0])\n",
        "    return X, target\n",
        "\n",
        "  def label2imgid(self):\n",
        "    self.return_imgid = 1\n",
        "\n",
        "  def __len__(self) -\u003e int:\n",
        "    return len(self.filename)\n",
        "\n",
        "  def extra_repr(self) -\u003e str:\n",
        "    lines = [\"Target type: {target_type}\", \"Split: {split}\"]\n",
        "    return '\\n'.join(lines).format(**self.__dict__)\n",
        "\n",
        "\n",
        "\"\"\"Here is the model.\"\"\"\n",
        "\n",
        "\n",
        "class GRBM(nn.Module):\n",
        "  \"\"\" Gaussian-Bernoulli Restricted Boltzmann Machines (GRBM) \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               visible_size,\n",
        "               hidden_size,\n",
        "               CD_step=1,\n",
        "               CD_burnin=0,\n",
        "               init_var=1e-0,\n",
        "               neg_MC_init_var=1e-0,\n",
        "               inference_method='Gibbs',\n",
        "               Langevin_step=10,\n",
        "               Langevin_eta=1.0,\n",
        "               is_anneal_Langevin=True,\n",
        "               Langevin_adjust_step=0) -\u003e None:\n",
        "    super().__init__()\n",
        "    # we use samples in [CD_burnin, CD_step) steps\n",
        "    assert CD_burnin \u003e= 0 and CD_burnin \u003c= CD_step\n",
        "    assert inference_method in ['Gibbs', 'Langevin', 'Langevin-Gibbs']\n",
        "\n",
        "    self.visible_size = visible_size\n",
        "    self.hidden_size = hidden_size\n",
        "    self.CD_step = CD_step\n",
        "    self.CD_burnin = CD_burnin\n",
        "    self.init_var = init_var\n",
        "    self.neg_MC_init_var = neg_MC_init_var\n",
        "    self.inference_method = inference_method\n",
        "    self.Langevin_step = Langevin_step\n",
        "    self.Langevin_eta = Langevin_eta\n",
        "    self.is_anneal_Langevin = is_anneal_Langevin\n",
        "    self.Langevin_adjust_step = Langevin_adjust_step\n",
        "\n",
        "    self.W = nn.Parameter(torch.Tensor(visible_size, hidden_size))\n",
        "    self.b = nn.Parameter(torch.Tensor(hidden_size))\n",
        "    self.mu = nn.Parameter(torch.Tensor(visible_size))\n",
        "    self.log_var = nn.Parameter(torch.Tensor(visible_size))\n",
        "    self.reset_parameters()\n",
        "\n",
        "  def reset_parameters(self):\n",
        "    nn.init.normal_(self.W,\n",
        "                    std=1.0 * self.init_var /\n",
        "                    np.sqrt(self.visible_size + self.hidden_size))\n",
        "    nn.init.constant_(self.b, 0.0)\n",
        "    nn.init.constant_(self.mu, 0.0)\n",
        "    nn.init.constant_(self.log_var,\n",
        "                      np.log(self.init_var))  # init variance = 1.0\n",
        "\n",
        "  def get_var(self):\n",
        "    return self.log_var.exp().clip(min=1e-8)\n",
        "\n",
        "  def set_Langevin_adjust_step(self, step):\n",
        "    self.Langevin_adjust_step = step\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def energy(self, v, h):\n",
        "    var = self.get_var()\n",
        "    eng = 0.5 * ((v - self.mu)**2 / var).sum(dim=1)\n",
        "    eng -= ((v / var).mm(self.W) * h).sum(dim=1) + h.mv(self.b)\n",
        "    return eng\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def marginal_energy(self, v):\n",
        "    var = self.get_var()\n",
        "    eng = 0.5 * ((v - self.mu)**2 / var).sum(dim=1)\n",
        "    eng -= F.softplus((v / var).mm(self.W) + self.b).sum(dim=1)\n",
        "    return eng\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def energy_grad_v(self, v, h):\n",
        "    B = v.shape[0]\n",
        "    var = self.get_var()\n",
        "    return ((v - self.mu) / var - h.mm(self.W.T) / var) / B\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def marginal_energy_grad_v(self, v):\n",
        "    B = v.shape[0]\n",
        "    var = self.get_var()\n",
        "    return ((v - self.mu) / var -\n",
        "            torch.sigmoid((v / var).mm(self.W) + self.b).mm(self.W.T) / var) / B\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def energy_grad_param(self, v, h):\n",
        "    var = self.get_var()\n",
        "    grad = {}\n",
        "    grad['W'] = -torch.einsum(\"bi,bj-\u003eij\", v / var, h) / v.shape[0]\n",
        "    grad['b'] = -h.mean(dim=0)\n",
        "    grad['mu'] = ((self.mu - v) / var).mean(dim=0)\n",
        "    grad['log_var'] = (-0.5 * (v - self.mu)**2 / var +\n",
        "                       ((v / var) * h.mm(self.W.T))).mean(dim=0)\n",
        "    return grad\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def marginal_energy_grad_param(self, v):\n",
        "    var = self.get_var()\n",
        "    vv = v / var\n",
        "    tmp = torch.sigmoid(vv.mm(self.W) + self.b)\n",
        "    grad = {}\n",
        "    grad['W'] = -torch.einsum(\"bi,bj-\u003eij\", vv, tmp) / v.shape[0]\n",
        "    grad['b'] = -tmp.mean(dim=0)\n",
        "    grad['mu'] = ((self.mu - v) / var).mean(dim=0)\n",
        "    grad['log_var'] = (-0.5 * (v - self.mu)**2 / var +\n",
        "                       (vv * tmp.mm(self.W.T))).mean(dim=0)\n",
        "    return grad\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def prob_h_given_v(self, v, var):\n",
        "    return torch.sigmoid((v / var).mm(self.W) + self.b)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def prob_v_given_h(self, h):\n",
        "    return h.mm(self.W.T) + self.mu\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def log_metropolis_ratio_v_given_h(self, v_old, v_new, h, grad_old, eta):\n",
        "    \"\"\" Metropolis-Hasting ratio of accepting the move from old to new state \"\"\"\n",
        "    eng_diff = -self.energy(v_new, h) + self.energy(v_old, h)\n",
        "    proposal_eng_new = - \\\n",
        "        torch.pow(v_old - v_new + eta *\n",
        "                  self.energy_grad_v(v_new, h), 2.0).sum(dim=1) / (4 * eta)\n",
        "    proposal_eng_old = - \\\n",
        "        torch.pow(v_new - v_old + eta * grad_old,\n",
        "                  2.0).sum(dim=1) / (4 * eta)\n",
        "\n",
        "    return eng_diff + proposal_eng_new - proposal_eng_old\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def log_metropolis_ratio_Gibbs_joint(self, v_old, h_old, v_new, h_new, var):\n",
        "    \"\"\" Metropolis-Hasting ratio of accepting the move from old to new state \"\"\"\n",
        "    marginal_eng_diff = -self.marginal_energy(v_new) + self.marginal_energy(\n",
        "        v_old)\n",
        "    eng_diff = -0.5 * (\n",
        "        (v_old - self.mu - h_new.mm(self.W.T))**2 / var).sum(dim=1) + 0.5 * (\n",
        "            (v_new - self.mu - h_old.mm(self.W.T))**2 / var).sum(dim=1)\n",
        "    return eng_diff + marginal_eng_diff\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def log_metropolis_ratio_Langevin_Gibbs_joint(self, v_old, h_old, v_new,\n",
        "                                                h_new, eta):\n",
        "    \"\"\" Metropolis-Hasting ratio of accepting the move from old to new state \"\"\"\n",
        "    eng_diff = -self.marginal_energy(v_new) + self.marginal_energy(v_old)\n",
        "    proposal_eng_new = - \\\n",
        "        torch.pow(v_old - v_new + eta * self.energy_grad_v(v_new, h_new), 2.0).sum(dim=1) / (4 * eta)\n",
        "    proposal_eng_old = - \\\n",
        "        torch.pow(v_new - v_old + eta * self.energy_grad_v(v_old, h_old), 2.0).sum(dim=1) / (4 * eta)\n",
        "\n",
        "    return eng_diff + proposal_eng_new - proposal_eng_old\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def log_metropolis_ratio_marginal(self, v_old, v_new, grad_old, eta):\n",
        "    \"\"\" Metropolis-Hasting ratio of accepting the move from old to new state \"\"\"\n",
        "    eng_diff = -self.marginal_energy(v_new) + self.marginal_energy(v_old)\n",
        "    proposal_eng_new = - \\\n",
        "        torch.pow(v_old - v_new + eta *\n",
        "                  self.marginal_energy_grad_v(v_new), 2.0).sum(dim=1) / (4 * eta)\n",
        "    proposal_eng_old = - \\\n",
        "        torch.pow(v_new - v_old + eta * grad_old,\n",
        "                  2.0).sum(dim=1) / (4 * eta)\n",
        "\n",
        "    return eng_diff + proposal_eng_new - proposal_eng_old\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def Gibbs_sampling_vh(self, v, num_steps=10, burn_in=0):\n",
        "    samples, var = [], self.get_var()\n",
        "    std = var.sqrt()\n",
        "    h = torch.bernoulli(self.prob_h_given_v(v, var))\n",
        "    for ii in range(num_steps):\n",
        "      # backward sampling\n",
        "      mu = self.prob_v_given_h(h)\n",
        "      v = mu + torch.randn_like(mu) * std\n",
        "\n",
        "      # forward sampling\n",
        "      h = torch.bernoulli(self.prob_h_given_v(v, var))\n",
        "\n",
        "      if ii \u003e= burn_in:\n",
        "        samples += [(v, h)]\n",
        "\n",
        "    return samples\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def Langevin_sampling_v(self,\n",
        "                          v,\n",
        "                          num_steps=10,\n",
        "                          eta=1.0e+0,\n",
        "                          burn_in=0,\n",
        "                          is_anneal=True,\n",
        "                          adjust_step=0):\n",
        "    eta_list = cosine_schedule(eta_max=eta, T=num_steps)\n",
        "    samples = []\n",
        "\n",
        "    for ii in range(num_steps):\n",
        "      eta_ii = eta_list[ii] if is_anneal else eta\n",
        "      grad_v = self.marginal_energy_grad_v(v)\n",
        "\n",
        "      v_new = v - eta_ii * grad_v + \\\n",
        "          torch.randn_like(v) * np.sqrt(eta_ii * 2)\n",
        "\n",
        "      if ii \u003e= adjust_step:\n",
        "        tmp_u = torch.rand(v.shape[0]).to(v.device)\n",
        "        log_ratio = self.log_metropolis_ratio_marginal(v, v_new, grad_v, eta_ii)\n",
        "        ratio = torch.minimum(torch.ones_like(log_ratio), log_ratio.exp())\n",
        "        v = v_new * (tmp_u \u003c ratio).float().view(\n",
        "            -1, 1) + v * (tmp_u \u003e= ratio).float().view(-1, 1)\n",
        "      else:\n",
        "        v = v_new\n",
        "\n",
        "      if ii \u003e= burn_in:\n",
        "        samples += [v]\n",
        "\n",
        "    return samples\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def Langevin_Gibbs_sampling_vh(self,\n",
        "                                 v,\n",
        "                                 num_steps=10,\n",
        "                                 num_steps_Langevin=10,\n",
        "                                 eta=1.0e+0,\n",
        "                                 burn_in=0,\n",
        "                                 is_anneal=True,\n",
        "                                 adjust_step=0):\n",
        "    samples, var = [], self.get_var()\n",
        "    eta_list = cosine_schedule(eta_max=eta, T=num_steps_Langevin)\n",
        "\n",
        "    h = torch.bernoulli(self.prob_h_given_v(v, var))\n",
        "    for ii in range(num_steps):\n",
        "      v_old, h_old = v, h\n",
        "      # backward sampling\n",
        "      for jj in range(num_steps_Langevin):\n",
        "        eta_jj = eta_list[jj] if is_anneal else eta\n",
        "        grad_v = self.energy_grad_v(v, h)\n",
        "        v = v - eta_jj * grad_v + torch.randn_like(v) * np.sqrt(eta_jj * 2)\n",
        "\n",
        "      # forward sampling\n",
        "      h = torch.bernoulli(self.prob_h_given_v(v, var))\n",
        "\n",
        "      if ii \u003e= adjust_step:\n",
        "        tmp_u = torch.rand(v.shape[0]).to(v.device)\n",
        "        log_ratio = self.log_metropolis_ratio_Langevin_Gibbs_joint(\n",
        "            v_old, h_old, v, h, eta)\n",
        "        ratio = torch.minimum(torch.ones_like(log_ratio), log_ratio.exp())\n",
        "        v = v * (tmp_u \u003c ratio).float().view(\n",
        "            -1, 1) + v_old * (tmp_u \u003e= ratio).float().view(-1, 1)\n",
        "        h = h * (tmp_u \u003c ratio).float().view(\n",
        "            -1, 1) + h_old * (tmp_u \u003e= ratio).float().view(-1, 1)\n",
        "\n",
        "      if ii \u003e= burn_in:\n",
        "        samples += [(v, h)]\n",
        "\n",
        "    return samples\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def reconstruction(self, v):\n",
        "    v, var = v.view(v.shape[0], -1), self.get_var()\n",
        "    prob_h = self.prob_h_given_v(v, var)\n",
        "    v_bar = self.prob_v_given_h(prob_h)\n",
        "    return F.mse_loss(v, v_bar)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def sampling(self, v_init, num_steps=1, save_gap=1):\n",
        "    v_shape = v_init.shape\n",
        "    v = v_init.view(v_shape[0], -1)\n",
        "    var = self.get_var()\n",
        "    var_mean = var.mean().item()\n",
        "\n",
        "    if self.inference_method == 'Gibbs':\n",
        "      samples = self.Gibbs_sampling_vh(v, num_steps=num_steps - 1)\n",
        "      samples = [xx[0] for xx in samples]  # extract v\n",
        "    elif self.inference_method == 'Langevin':\n",
        "      samples = self.Langevin_sampling_v(v,\n",
        "                                         num_steps=num_steps - 1,\n",
        "                                         eta=self.Langevin_eta * var_mean,\n",
        "                                         is_anneal=self.is_anneal_Langevin,\n",
        "                                         adjust_step=self.Langevin_adjust_step)\n",
        "    elif self.inference_method == 'Langevin-Gibbs':\n",
        "      samples = self.Langevin_Gibbs_sampling_vh(\n",
        "          v,\n",
        "          num_steps=num_steps - 1,\n",
        "          num_steps_Langevin=self.Langevin_step,\n",
        "          eta=self.Langevin_eta * var_mean,\n",
        "          is_anneal=self.is_anneal_Langevin,\n",
        "          adjust_step=self.Langevin_adjust_step)\n",
        "      samples = [xx[0] for xx in samples]  # extract v\n",
        "\n",
        "    # use conditional mean as the last sample\n",
        "    h = torch.bernoulli(self.prob_h_given_v(samples[-1], var))\n",
        "    mu = self.prob_v_given_h(h)\n",
        "    v_list = [(0, v_init)] + [(ii + 1, samples[ii].view(v_shape).detach())\n",
        "                              for ii in range(num_steps - 1)\n",
        "                              if (ii + 1) % save_gap == 0\n",
        "                             ] + [(num_steps, mu.view(v_shape).detach())]\n",
        "\n",
        "    return v_list\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def positive_grad(self, v):\n",
        "    h = torch.bernoulli(self.prob_h_given_v(v, self.get_var()))\n",
        "    grad = self.energy_grad_param(v, h)\n",
        "    return grad\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def negative_grad(self, v):\n",
        "    var = self.get_var()\n",
        "    var_mean = var.mean().item()\n",
        "    if self.inference_method == 'Gibbs':\n",
        "      samples = self.Gibbs_sampling_vh(v,\n",
        "                                       num_steps=self.CD_step,\n",
        "                                       burn_in=self.CD_burnin)\n",
        "      v_neg = torch.cat([xx[0] for xx in samples], dim=0)\n",
        "      h_neg = torch.cat([xx[1] for xx in samples], dim=0)\n",
        "      grad = self.energy_grad_param(v_neg, h_neg)\n",
        "    elif self.inference_method == 'Langevin':\n",
        "      samples = self.Langevin_sampling_v(v,\n",
        "                                         num_steps=self.CD_step,\n",
        "                                         burn_in=self.CD_burnin,\n",
        "                                         eta=self.Langevin_eta * var_mean,\n",
        "                                         is_anneal=self.is_anneal_Langevin,\n",
        "                                         adjust_step=self.Langevin_adjust_step)\n",
        "      v_neg = torch.cat(samples, dim=0)\n",
        "      grad = self.marginal_energy_grad_param(v_neg)\n",
        "\n",
        "    elif self.inference_method == 'Langevin-Gibbs':\n",
        "      samples = self.Langevin_Gibbs_sampling_vh(\n",
        "          v,\n",
        "          num_steps=self.CD_step,\n",
        "          burn_in=self.CD_burnin,\n",
        "          num_steps_Langevin=self.Langevin_step,\n",
        "          eta=self.Langevin_eta * var_mean,\n",
        "          is_anneal=self.is_anneal_Langevin,\n",
        "          adjust_step=self.Langevin_adjust_step)\n",
        "      v_neg = torch.cat([xx[0] for xx in samples], dim=0)\n",
        "      h_neg = torch.cat([xx[1] for xx in samples], dim=0)\n",
        "      grad = self.energy_grad_param(v_neg, h_neg)\n",
        "\n",
        "    return grad\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def CD_grad(self, v):\n",
        "    v = v.view(v.shape[0], -1)\n",
        "    # postive gradient\n",
        "    grad_pos = self.positive_grad(v)\n",
        "\n",
        "    # negative gradient\n",
        "    v_neg = torch.randn_like(v) * self.neg_MC_init_var\n",
        "    grad_neg = self.negative_grad(v_neg)\n",
        "\n",
        "    # compute update\n",
        "    for name, param in self.named_parameters():\n",
        "      param.grad = grad_pos[name] - grad_neg[name]\n",
        "\n",
        "\n",
        "\"\"\"Here is the config of running experiments.\"\"\"\n",
        "\n",
        "\n",
        "def get_config(pid):\n",
        "  config = {}\n",
        "  config['dataset'] = 'GMM'\n",
        "  # config['dataset'] = 'MNIST'\n",
        "  # config['dataset'] = 'CIFAR10'\n",
        "  # config['dataset'] = 'CelebA'\n",
        "  config['cuda'] = True\n",
        "  config['model'] = 'GRBM'\n",
        "  config['batch_size'] = 128\n",
        "  config['epochs'] = 10000\n",
        "  config['lr'] = 1.0e-2\n",
        "  config['clip_norm'] = 10.0\n",
        "  config['wd'] = 0.0e-4\n",
        "  config['seed'] = 1  # random seed\n",
        "  config['optimizer'] = 'SGD'\n",
        "  config['resume'] = 0\n",
        "  config[\n",
        "      'is_vis_verbose'] = True  # visualize sampling process, filters, hiddens if True\n",
        "  config['init_var'] = 1e-0  # init variance of GRBM\n",
        "  config['hidden_size'] = 4096\n",
        "  config['CD_step'] = 100\n",
        "  config['CD_burnin'] = 0\n",
        "  config['Langevin_step'] = 1\n",
        "  config['Langevin_eta'] = 1.0e+1\n",
        "  config['is_anneal_Langevin'] = True\n",
        "  config['Langevin_adjust_step'] = 100\n",
        "  config['use_replay_buffer'] = False  # ignore for now\n",
        "  config['is_anneal_data_noise'] = False  # ignore for now\n",
        "  config['inference_method'] = 'Gibbs'\n",
        "  # config['inference_method'] = 'Langevin'\n",
        "  # config['inference_method'] = 'Langevin-Gibbs'\n",
        "  config['neg_MC_init_var'] = 1.0  # ignore for now\n",
        "  config['data_init_var'] = 2.0  # ignore for now\n",
        "  config['sampling_batch_size'] = 100\n",
        "  config['sampling_steps'] = config['CD_step']\n",
        "  config['sampling_gap'] = min(5, config['sampling_steps'])\n",
        "  config['sampling_nrow'] = 10\n",
        "\n",
        "  if config['dataset'] == 'GMM':\n",
        "    config['batch_size'] = 100\n",
        "    config['num_samples'] = 1000\n",
        "    config['height'] = 1\n",
        "    config['width'] = 1\n",
        "    config['channel'] = 2\n",
        "    config['log_interval'] = 100\n",
        "    config['save_interval'] = 50000\n",
        "    config['epochs'] = 50000\n",
        "    config['hidden_size'] = 256\n",
        "  elif config['dataset'] == 'MNIST':\n",
        "    config['height'] = 28\n",
        "    config['width'] = 28\n",
        "    config['channel'] = 1\n",
        "    config['img_mean'] = torch.tensor([0.1307])\n",
        "    config['img_std'] = torch.tensor([0.3081])\n",
        "    config['log_interval'] = 10\n",
        "    config['save_interval'] = 100\n",
        "  elif config['dataset'] == 'CIFAR10':\n",
        "    config['height'] = 32\n",
        "    config['width'] = 32\n",
        "    config['channel'] = 3\n",
        "    config['img_mean'] = torch.tensor([0.4914, 0.4822, 0.4465])\n",
        "    config['img_std'] = torch.tensor([0.2470, 0.2435, 0.2616])\n",
        "    config['log_interval'] = 10\n",
        "    config['save_interval'] = 10000\n",
        "    # config['log_interval'] = 1\n",
        "    # config['save_interval'] = 10\n",
        "  elif config['dataset'] == 'CelebA':\n",
        "    config['height'] = 32\n",
        "    config['width'] = 32\n",
        "    config['channel'] = 3\n",
        "    config['crop_size'] = 140\n",
        "    config['img_mean'] = torch.tensor([0.5240, 0.4152, 0.3590])\n",
        "    config['img_std'] = torch.tensor([0.2868, 0.2530, 0.2453])\n",
        "    config['log_interval'] = 1\n",
        "    config['save_interval'] = 5\n",
        "\n",
        "  config[\n",
        "      'visible_size'] = config['height'] * config['width'] * config['channel']\n",
        "\n",
        "  if config['inference_method'] == 'Gibbs':\n",
        "    config[\n",
        "        'exp_folder'] = f\"exp/{config['dataset']}_{config['model']}_{pid}_inference={config['inference_method']}_H={config['hidden_size']}_lr={config['lr']}_B={config['batch_size']}_initvar={config['init_var']}_CD={config['CD_step']}_burnin={config['CD_burnin']}_is_anneal_data_noise={config['is_anneal_data_noise']}_{config['optimizer']}\"\n",
        "  elif config['inference_method'] == 'Langevin':\n",
        "    config[\n",
        "        'exp_folder'] = f\"exp/{config['dataset']}_{config['model']}_{pid}_inference={config['inference_method']}_is_anneal_Langevin={config['is_anneal_Langevin']}_Langevin_adjust_step={config['Langevin_adjust_step']}_Langevin_eta={config['Langevin_eta']}_H={config['hidden_size']}_lr={config['lr']}_B={config['batch_size']}_initvar={config['init_var']}_CD={config['CD_step']}_burnin={config['CD_burnin']}_is_anneal_data_noise={config['is_anneal_data_noise']}_{config['optimizer']}\"\n",
        "  elif config['inference_method'] == 'Langevin-Gibbs':\n",
        "    config[\n",
        "        'exp_folder'] = f\"exp/{config['dataset']}_{config['model']}_{pid}_inference={config['inference_method']}_is_anneal_Langevin={config['is_anneal_Langevin']}_Langevin_adjust_step={config['Langevin_adjust_step']}_Langevin_step={config['Langevin_step']}_Langevin_eta={config['Langevin_eta']}_H={config['hidden_size']}_lr={config['lr']}_B={config['batch_size']}_initvar={config['init_var']}_CD={config['CD_step']}_burnin={config['CD_burnin']}_is_anneal_data_noise={config['is_anneal_data_noise']}_{config['optimizer']}\"\n",
        "\n",
        "  return config\n",
        "\n",
        "\n",
        "\"\"\"Here is the main training loop.\"\"\"\n",
        "\n",
        "\n",
        "def train_model():\n",
        "  \"\"\"Let us train a GRBM and see how it performs\"\"\"\n",
        "  pid = os.getpid()\n",
        "  config = get_config(pid)\n",
        "  if not os.path.isdir(config['exp_folder']):\n",
        "    os.makedirs(config['exp_folder'])\n",
        "\n",
        "  log_file = os.path.join(config['exp_folder'], f'log_exp_{pid}.txt')\n",
        "  logger = setup_logging('INFO', log_file)\n",
        "  logger.info('Writing log file to {}'.format(log_file))\n",
        "\n",
        "  with open(os.path.join(config['exp_folder'], f'config_{pid}.json'),\n",
        "            'w') as outfile:\n",
        "    json.dump(config, outfile, cls=TensorEncoder, indent=4)\n",
        "\n",
        "  if config['dataset'] == 'GMM':\n",
        "    gmm_model = GMM(torch.tensor([0.33, 0.33, 0.34]),\n",
        "                    torch.tensor([[-5, -5], [5, -5], [0, 5]]),\n",
        "                    torch.tensor([[1.25, 0.5], [1.25, 0.5], [0.5,\n",
        "                                                             1.25]])).cuda()\n",
        "    # gmm_model = GMM(torch.tensor([0.33, 0.33, 0.34]),\n",
        "    #                 torch.tensor([[-5, -5], [5, -5], [0, 5]]),\n",
        "    #                 torch.tensor([[1, 1], [1, 1], [1, 1]])).cuda()\n",
        "    # gmm_model = GMM(torch.tensor([1.0]), torch.tensor([[2, 2]]),\n",
        "    #              torch.tensor([[2.0, 0.5]])).cuda()\n",
        "\n",
        "    vis_density_GMM(gmm_model, config)\n",
        "    samples = gmm_model.sampling(config['num_samples'])\n",
        "    vis_2D_samples(samples.cpu().numpy(), config, tags='ground_truth')\n",
        "    train_set = GMMDataset(samples)\n",
        "  elif config['dataset'] == 'MNIST':\n",
        "    train_set = datasets.MNIST('./data',\n",
        "                               train=True,\n",
        "                               download=True,\n",
        "                               transform=transforms.Compose([\n",
        "                                   transforms.ToTensor(),\n",
        "                                   transforms.Normalize(config['img_mean'],\n",
        "                                                        config['img_std'])\n",
        "                               ]))\n",
        "  elif config['dataset'] == 'CIFAR10':\n",
        "    train_set = datasets.CIFAR10('./data',\n",
        "                                 train=True,\n",
        "                                 download=True,\n",
        "                                 transform=transforms.Compose([\n",
        "                                     transforms.ToTensor(),\n",
        "                                     transforms.Normalize(\n",
        "                                         config['img_mean'], config['img_std'])\n",
        "                                 ]))\n",
        "  elif config['dataset'] == 'CelebA':\n",
        "    train_set = CelebA('./data',\n",
        "                       split='train',\n",
        "                       download=False,\n",
        "                       transform=transforms.Compose([\n",
        "                           transforms.CenterCrop(config['crop_size']),\n",
        "                           transforms.Resize(config['height']),\n",
        "                           transforms.ToTensor(),\n",
        "                           transforms.Normalize(config['img_mean'],\n",
        "                                                config['img_std'])\n",
        "                       ]))\n",
        "\n",
        "  # use a subset to debug\n",
        "  # train_set = torch.utils.data.Subset(train_set, range(128))\n",
        "\n",
        "  train_loader = torch.utils.data.DataLoader(train_set,\n",
        "                                             batch_size=config['batch_size'],\n",
        "                                             shuffle=True)\n",
        "\n",
        "  model = GRBM(config['visible_size'],\n",
        "               config['hidden_size'],\n",
        "               CD_step=config['CD_step'],\n",
        "               CD_burnin=config['CD_burnin'],\n",
        "               init_var=config['init_var'],\n",
        "               neg_MC_init_var=config['neg_MC_init_var'],\n",
        "               inference_method=config['inference_method'],\n",
        "               Langevin_step=config['Langevin_step'],\n",
        "               Langevin_eta=config['Langevin_eta'],\n",
        "               is_anneal_Langevin=config['is_anneal_Langevin'],\n",
        "               Langevin_adjust_step=config['Langevin_adjust_step'])\n",
        "\n",
        "  if config['cuda']:\n",
        "    model.cuda()\n",
        "\n",
        "  param_wd, param_no_wd = [], []\n",
        "  for xx, yy in model.named_parameters():\n",
        "    if 'W' in xx:\n",
        "      param_wd += [yy]\n",
        "    else:\n",
        "      param_no_wd += [yy]\n",
        "\n",
        "  if config['optimizer'] == 'Adam':\n",
        "    optimizer = optim.Adam([{\n",
        "        'params': param_no_wd,\n",
        "        'weight_decay': 0\n",
        "    }, {\n",
        "        'params': param_wd\n",
        "    }],\n",
        "                           lr=config['lr'],\n",
        "                           weight_decay=config['wd'])\n",
        "  else:\n",
        "    optimizer = optim.SGD([{\n",
        "        'params': param_no_wd,\n",
        "        'weight_decay': 0\n",
        "    }, {\n",
        "        'params': param_wd\n",
        "    }],\n",
        "                          lr=config['lr'],\n",
        "                          momentum=0.0,\n",
        "                          weight_decay=config['wd'])\n",
        "\n",
        "  scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
        "      optimizer, config['epochs'])\n",
        "\n",
        "  if config['resume'] \u003e 0:\n",
        "    load(model, config['exp_folder'], config['resume'])\n",
        "\n",
        "  stats_dict = defaultdict(list)\n",
        "  for epoch in range(config['resume']):\n",
        "    scheduler.step()\n",
        "\n",
        "  is_show_training_data = False\n",
        "  data_std = [\n",
        "      math.sqrt(config['data_init_var']) *\n",
        "      (1 + math.cos(tt * math.pi / config['epochs'])) / 2\n",
        "      for tt in range(config['epochs'])\n",
        "  ]\n",
        "  for epoch in range(config['resume'] + 1, config['epochs'] + 1):\n",
        "    recon_loss = train(model,\n",
        "                       train_loader,\n",
        "                       optimizer,\n",
        "                       config,\n",
        "                       is_anneal_data_noise=config['is_anneal_data_noise'],\n",
        "                       std=data_std[epoch - 1])\n",
        "\n",
        "    var = model.get_var().detach().cpu().numpy()\n",
        "    var_list = var.tolist()\n",
        "    log_var_list = np.log(var).tolist()\n",
        "    stats_dict['Variance'] += var_list\n",
        "    stats_dict['Log Variance'] += log_var_list\n",
        "    stats_dict['Epoch'] += [epoch] * len(var_list)\n",
        "    stats_dict['Units'] += range(1, len(var_list) + 1)\n",
        "    stats_dict['Reconstruction MSE Loss'] += [recon_loss] * len(var_list)\n",
        "\n",
        "    if epoch % config['log_interval'] == 0:\n",
        "      if config['dataset'] == 'GMM':\n",
        "        logger.info(\n",
        "            f'PID={pid} || {epoch} epoch || mean = {model.mu.detach().cpu().numpy()} || var={model.get_var().detach().cpu().numpy()} || Reconstruction Loss = {recon_loss}'\n",
        "        )\n",
        "      else:\n",
        "        logger.info(\n",
        "            f'PID={pid} || {epoch} epoch || var={model.get_var().mean().item()} || Reconstruction Loss = {recon_loss}'\n",
        "        )\n",
        "\n",
        "      visualize_sampling(model,\n",
        "                         epoch,\n",
        "                         config,\n",
        "                         is_show_gif=config['is_vis_verbose'])\n",
        "      if config['is_vis_verbose']:\n",
        "        # visualize filters\n",
        "        filters = model.W.T.view(model.W.shape[1], config['channel'],\n",
        "                                 config['height'], config['width'])\n",
        "        utils.save_image(\n",
        "            filters,\n",
        "            f\"{config['exp_folder']}/filters_epoch_{epoch:05d}.png\",\n",
        "            nrow=8,\n",
        "            normalize=True,\n",
        "            padding=1,\n",
        "            pad_value=1.0)\n",
        "\n",
        "        # visualize hidden states\n",
        "        data, _ = next(iter(train_loader))\n",
        "        h_pos = model.prob_h_given_v(\n",
        "            data.view(data.shape[0], -1).cuda(), model.get_var())\n",
        "        utils.save_image(h_pos.view(1, 1, -1, config['hidden_size']),\n",
        "                         f\"{config['exp_folder']}/hidden_epoch_{epoch:05d}.png\",\n",
        "                         normalize=True)\n",
        "\n",
        "        # visualize one mini-batch of training data\n",
        "        if not is_show_training_data and config['dataset'] != 'GMM':\n",
        "          data, _ = next(iter(train_loader))\n",
        "          mean = config['img_mean'].view(1, -1, 1, 1).to(data.device)\n",
        "          std = config['img_std'].view(1, -1, 1, 1).to(data.device)\n",
        "          vis_data = (data * std + mean).clamp(min=0, max=1)\n",
        "          utils.save_image(\n",
        "              utils.make_grid(vis_data,\n",
        "                              nrow=config['sampling_nrow'],\n",
        "                              normalize=False,\n",
        "                              padding=1,\n",
        "                              pad_value=1.0).cpu(),\n",
        "              f\"{config['exp_folder']}/training_imgs.png\")\n",
        "          is_show_training_data = True\n",
        "\n",
        "    # save samples periodically.\n",
        "    if epoch % config['save_interval'] == 0:\n",
        "      save(model, config['exp_folder'], epoch)\n",
        "      # plotting a lot of points is slow so we use another thread\n",
        "      thread = Thread(target=plot_curves, args=(stats_dict, config))\n",
        "      thread.start()\n",
        "    #   thread.join()\n",
        "\n",
        "    scheduler.step()\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "  train_model()\n",
        "  # test_model()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkBWeFQHw0v8"
      },
      "source": [
        "# Evaluate Trained Models\n",
        "\n",
        "Required packages:\n",
        "\n",
        "*   pytorch_fid\n",
        "*   seaborn\n",
        "*   pandas\n",
        "*   matplotlib\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wvzKnPrTyoAo"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import pdb\n",
        "import glob\n",
        "import json\n",
        "import subprocess\n",
        "from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser\n",
        "import imageio\n",
        "\n",
        "import numpy as np\n",
        "from tqdm import tqdm\n",
        "import subprocess\n",
        "import seaborn as sns\n",
        "import pandas as pd\n",
        "from collections import defaultdict\n",
        "import matplotlib\n",
        "\n",
        "matplotlib.use('Agg')\n",
        "import matplotlib.pyplot as plt  # NOQA\n",
        "\n",
        "sns.set_theme(style=\"darkgrid\")\n",
        "\n",
        "from scipy.stats import entropy\n",
        "\n",
        "import torch\n",
        "from torchvision import datasets, transforms, utils\n",
        "from torch import nn\n",
        "from torch.nn import functional as F\n",
        "import torch.utils.data\n",
        "from torchvision.models.inception import inception_v3\n",
        "from model.grbm_minimal import CelebA, GRBM\n",
        "\n",
        "parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)\n",
        "parser.add_argument('--path',\n",
        "                    type=str,\n",
        "                    default=None,\n",
        "                    help='path of model folder')\n",
        "parser.add_argument('--is_test',\n",
        "                    type=bool,\n",
        "                    default=True,\n",
        "                    help='evaluate using training or testing set')\n",
        "parser.add_argument('--gpu_list', nargs='+', default=['0', '1'], help='gpu ID')\n",
        "\n",
        "\n",
        "class CustomImageDataset(torch.utils.data.Dataset):\n",
        "\n",
        "  def __init__(self, root_dir, transform=None):\n",
        "    \"\"\"\n",
        "        config:        \n",
        "            root_dir (string): Directory with all the images.\n",
        "            transform (callable, optional): Optional transform to be applied\n",
        "                on a sample.\n",
        "        \"\"\"\n",
        "\n",
        "    self.root_dir = root_dir\n",
        "    self.img_names = glob.glob(os.path.join(self.root_dir, '*.png'))\n",
        "    self.transform = transform\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.img_names)\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    image = imageio.imread(self.img_names[idx])\n",
        "\n",
        "    if self.transform:\n",
        "      sample = self.transform(image)\n",
        "      if sample.shape[0] == 1:\n",
        "        sample = sample.repeat((3, 1, 1))\n",
        "\n",
        "    return sample\n",
        "\n",
        "\n",
        "def inception_score(img_path, cuda=True, batch_size=32, resize=False, splits=1):\n",
        "  \"\"\"Computes the inception score of the generated images imgs\n",
        "\n",
        "    imgs -- Torch dataset of (3xHxW) numpy images normalized in the range [-1, 1]\n",
        "    cuda -- whether or not to run on GPU\n",
        "    batch_size -- batch size for feeding into Inception v3\n",
        "    splits -- number of splits\n",
        "    \"\"\"\n",
        "  custom_transform = transforms.Compose(\n",
        "      [transforms.ToTensor(),\n",
        "       transforms.Normalize((0.5,), (0.5,))])\n",
        "  dataset = CustomImageDataset(img_path, transform=custom_transform)\n",
        "\n",
        "  N = len(dataset)\n",
        "\n",
        "  assert batch_size \u003e 0\n",
        "  assert N \u003e batch_size\n",
        "\n",
        "  # Set up dtype\n",
        "  if cuda:\n",
        "    dtype = torch.cuda.FloatTensor\n",
        "  else:\n",
        "    if torch.cuda.is_available():\n",
        "      print(\n",
        "          \"WARNING: You have a CUDA device, so you should probably set cuda=True\"\n",
        "      )\n",
        "    dtype = torch.FloatTensor\n",
        "\n",
        "  # Set up dataloader\n",
        "  dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)\n",
        "\n",
        "  # Load inception model\n",
        "  inception_model = inception_v3(pretrained=True,\n",
        "                                 transform_input=False).type(dtype)\n",
        "  inception_model.eval()\n",
        "  up = nn.Upsample(size=(299, 299), mode='bilinear').type(dtype)\n",
        "\n",
        "  def get_pred(x):\n",
        "    if resize:\n",
        "      x = up(x)\n",
        "    x = inception_model(x)\n",
        "    return F.softmax(x, dim=1).data.cpu().numpy()\n",
        "\n",
        "  # Get predictions\n",
        "  preds = np.zeros((N, 1000))\n",
        "\n",
        "  for i, batch in enumerate(dataloader, 0):\n",
        "    # batch = batch.type(dtype)\n",
        "    # batchv = Variable(batch)\n",
        "    batchv = batch.type(dtype)\n",
        "    batch_size_i = batch.size()[0]\n",
        "\n",
        "    preds[i * batch_size:i * batch_size + batch_size_i] = get_pred(batchv)\n",
        "\n",
        "  # Now compute the mean kl-div\n",
        "  split_scores = []\n",
        "\n",
        "  for k in range(splits):\n",
        "    part = preds[k * (N // splits):(k + 1) * (N // splits), :]\n",
        "    py = np.mean(part, axis=0)\n",
        "    scores = []\n",
        "    for i in range(part.shape[0]):\n",
        "      pyx = part[i, :]\n",
        "      scores.append(entropy(pyx, py))\n",
        "    split_scores.append(np.exp(np.mean(scores)))\n",
        "\n",
        "  return np.mean(split_scores), np.std(split_scores)\n",
        "\n",
        "\n",
        "def load(model, model_name):\n",
        "  data = torch.load(model_name)\n",
        "  model.load_state_dict(data['model'])\n",
        "\n",
        "\n",
        "def eval_model():\n",
        "  input_args = parser.parse_args()\n",
        "  config_file = glob.glob(os.path.join(input_args.path, 'config*.json'))[0]\n",
        "  config = json.load(open(config_file, 'r'))\n",
        "  config['img_mean'] = torch.from_numpy(np.array(config['img_mean']))\n",
        "  config['img_std'] = torch.from_numpy(np.array(config['img_std']))\n",
        "  config['exp_folder'] = input_args.path\n",
        "  config['cuda'] = True\n",
        "  config['split'] = 'test' if input_args.is_test else 'train'\n",
        "  config['batch_size'] = 100\n",
        "  config['sampling_steps'] = 100\n",
        "  config['sampling_gap'] = 5\n",
        "  config['sampling_nrow'] = 10\n",
        "  config['model_img_folder'] = f\"{config['exp_folder']}/gen_imgs\"\n",
        "\n",
        "  if config['dataset'] == 'MNIST':\n",
        "    data_set = datasets.MNIST('./data',\n",
        "                              train=True,\n",
        "                              download=True,\n",
        "                              transform=transforms.Compose(\n",
        "                                  [transforms.ToTensor()]))\n",
        "    config['gt_img_folder'] = f\"data/MNIST/imgs_{config['split']}\"\n",
        "  elif config['dataset'] == 'CIFAR10':\n",
        "    data_set = datasets.CIFAR10('./data',\n",
        "                                train=True,\n",
        "                                download=True,\n",
        "                                transform=transforms.Compose(\n",
        "                                    [transforms.ToTensor()]))\n",
        "    config['gt_img_folder'] = f\"data/CIFAR10/imgs_{config['split']}\"\n",
        "  elif config['dataset'] == 'CelebA':\n",
        "    data_set = CelebA('./data',\n",
        "                      split='train',\n",
        "                      download=False,\n",
        "                      transform=transforms.Compose([\n",
        "                          transforms.CenterCrop(config['crop_size']),\n",
        "                          transforms.Resize(config['height']),\n",
        "                          transforms.ToTensor()\n",
        "                      ]))\n",
        "    config['gt_img_folder'] = f\"data/celeba/imgs_{config['split']}\"\n",
        "\n",
        "  if not os.path.isdir(config['gt_img_folder']):\n",
        "    os.makedirs(config['gt_img_folder'])\n",
        "\n",
        "  if not os.path.isdir(config['model_img_folder']):\n",
        "    os.makedirs(config['model_img_folder'])\n",
        "\n",
        "  data_set = torch.utils.data.Subset(data_set, range(128))\n",
        "\n",
        "  data_loader = torch.utils.data.DataLoader(data_set,\n",
        "                                            batch_size=config['batch_size'])\n",
        "\n",
        "  model = GRBM(config['visible_size'],\n",
        "               config['hidden_size'],\n",
        "               CD_step=config['CD_step'],\n",
        "               CD_burnin=config['CD_burnin'],\n",
        "               init_var=config['init_var'],\n",
        "               neg_MC_init_var=config['neg_MC_init_var'],\n",
        "               inference_method=config['inference_method'],\n",
        "               Langevin_step=config['Langevin_step'],\n",
        "               Langevin_eta=config['Langevin_eta'],\n",
        "               is_anneal_Langevin=config['is_anneal_Langevin'],\n",
        "               Langevin_adjust_step=config['Langevin_adjust_step'])\n",
        "\n",
        "  if config['cuda']:\n",
        "    model.cuda()\n",
        "\n",
        "  model_names = glob.glob(os.path.join(config['exp_folder'], '*.pt'))\n",
        "\n",
        "  score_dict = defaultdict(list)\n",
        "  best_epoch, best_FID_score, best_IS_score = 0, 1e+5, 0\n",
        "  if isinstance(config['img_mean'], torch.Tensor):\n",
        "    config['img_mean'] = config['img_mean'].view(1, -1, 1, 1).cuda()\n",
        "    config['img_std'] = config['img_std'].view(1, -1, 1, 1).cuda()\n",
        "\n",
        "  for name in model_names:\n",
        "    load(model, name)\n",
        "    count = 0\n",
        "    for data, _ in tqdm(data_loader):\n",
        "      imgs = (data.permute(0, 2, 3, 1).numpy() * 255.0).astype(np.uint8)\n",
        "      for ii in range(data.shape[0]):\n",
        "        imageio.imsave(f\"{config['gt_img_folder']}/img_{count+ii+1:06d}.png\",\n",
        "                       imgs[ii])\n",
        "\n",
        "      v_init = torch.randn(config['batch_size'], config['channel'],\n",
        "                           config['height'],\n",
        "                           config['width']).cuda() * config['neg_MC_init_var']\n",
        "      v_list = model.sampling(v_init,\n",
        "                              num_steps=config['sampling_steps'],\n",
        "                              save_gap=config['sampling_gap'])\n",
        "\n",
        "      imgs = ((v_list[-1][1] * config['img_std'] + config['img_mean']).clamp(\n",
        "          min=0, max=1).cpu().permute(0, 2, 3, 1).numpy() * 255.0).astype(\n",
        "              np.uint8)\n",
        "\n",
        "      for ii in range(data.shape[0]):\n",
        "        imageio.imsave(f\"{config['model_img_folder']}/img_{count+ii+1:06d}.png\",\n",
        "                       imgs[ii])\n",
        "\n",
        "      count += data.shape[0]\n",
        "\n",
        "    cmd = f\"/opt/conda/bin/python -m pytorch_fid {config['model_img_folder']} {config['gt_img_folder']} --device cuda:0\"\n",
        "\n",
        "    process = subprocess.Popen(cmd,\n",
        "                               stdout=subprocess.PIPE,\n",
        "                               stderr=None,\n",
        "                               shell=True)\n",
        "    output = process.communicate()[0]\n",
        "    epoch = name.split('/')[-1].split('-')[1][:-3]\n",
        "    FID_score = float(output.decode('UTF-8').split(':')[1].strip())\n",
        "\n",
        "    IS_score_mean, IS_score_std = inception_score(config['model_img_folder'],\n",
        "                                                  cuda=True,\n",
        "                                                  batch_size=32,\n",
        "                                                  resize=True,\n",
        "                                                  splits=10)\n",
        "\n",
        "    score_dict['FID'] += [FID_score]\n",
        "    score_dict['Epoch'] += [int(epoch)]\n",
        "\n",
        "    print(\n",
        "        f'Epoch = {epoch} || FID score = {FID_score} || Inception Score = {IS_score_mean}'\n",
        "    )\n",
        "    if FID_score \u003c best_FID_score:\n",
        "      best_epoch = epoch\n",
        "      best_FID_score = FID_score\n",
        "      best_IS_score = IS_score_mean\n",
        "\n",
        "  print(\n",
        "      f'Best Epoch = {best_epoch} || Best FID score = {best_FID_score} || Best Inception score = {best_IS_score}'\n",
        "  )\n",
        "  print(config['exp_folder'])\n",
        "\n",
        "  df = pd.DataFrame(data=score_dict)\n",
        "  sns_plot = sns.lineplot(data=df, x=\"Epoch\", y=\"FID\", sort=True)\n",
        "  fig = sns_plot.get_figure()\n",
        "  fig.savefig(f\"{config['exp_folder']}/FID_curve.png\")\n",
        "  plt.clf()\n",
        "\n",
        "# ignored\n",
        "def parallel_eval_model():\n",
        "\n",
        "  def last_7chars(x):\n",
        "    return (x[-7:])\n",
        "\n",
        "  input_args = parser.parse_args()\n",
        "  model_names = sorted(glob.glob(os.path.join(input_args.path, 'epoch*')),\n",
        "                       key=last_7chars)\n",
        "  num_jobs_parallel = len(input_args.gpu_list)\n",
        "  results = {}\n",
        "\n",
        "  for ii in range(0, len(model_names), num_jobs_parallel):\n",
        "    process_list = []\n",
        "    for jj in range(min(num_jobs_parallel, len(model_names) - ii)):\n",
        "      name = model_names[ii + jj]\n",
        "      command = [\n",
        "          'python', '-m', 'pytorch_fid', name, 'data/MNIST/test', '--device',\n",
        "          'cuda:{}'.format(input_args.gpu_list[jj])\n",
        "      ]\n",
        "      process_list += [\n",
        "          subprocess.Popen(command, encoding=\"utf8\", stdout=subprocess.PIPE)\n",
        "      ]\n",
        "\n",
        "    for process in process_list:\n",
        "      process.wait()\n",
        "\n",
        "    for jj, process in enumerate(process_list):\n",
        "      epoch = model_names[ii + jj][-7:]\n",
        "      results[epoch] = float(process.stdout.readline().rstrip('\\n')[4:])\n",
        "\n",
        "  sorted_results = list(sorted(results.items(), key=lambda item: item[1]))\n",
        "  with open(os.path.join(input_args.path, 'test_results.txt'), 'w') as ff:\n",
        "    ff.write('{}\\n'.format(input_args.path))\n",
        "    for ii in range(len(sorted_results)):\n",
        "      ff.write('FID score @ Epoch = {} \u0026 {}\\n'.format(sorted_results[ii][1],\n",
        "                                                      sorted_results[ii][0]))\n",
        "\n",
        "  print('Best FID score \u0026 Epoch = {} \u0026 {}'.format(sorted_results[0][1],\n",
        "                                                  sorted_results[0][0]))\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "  eval_model()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "GRBM.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "141SyUTymzWH732EisKb1c_YH_7RMWmxx",
          "timestamp": 1661877437325
        },
        {
          "file_id": "1TkaSq-5gu2Cv6NPxsQ-il03J4pypwben",
          "timestamp": 1661877209930
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
